#include "Common.h"
#include "Shader.h"
#include "Utility.h"
#include "Render.h"
#include <iostream>
#include <vector>
#include <d3dx9.h>

Shader::Shader()
	: mVertexShader(NULL)
	, mPixelShader(NULL)
	, mVertexFactoryType(NULL)
{
}

Shader::~Shader()
{
	SAFE_RELEASE(mVertexShader);
	SAFE_RELEASE(mPixelShader);
}

bool Shader::Init()
{
	return true;
}

class D3DIncludeEnvironment : public ID3DXInclude
{
public:
	D3DIncludeEnvironment() {}
	STDMETHOD(Open)(D3DXINCLUDE_TYPE Type,LPCSTR Name,LPCVOID ParentData,LPCVOID* Data,UINT* Bytes)
	{
		std::string fileName(Name);
		std::string shaderSource;
		if (!Utility::LoadFileToString(SHADER_PATH + fileName, shaderSource))
			return D3DERR_NOTFOUND;

		char* fileContents = new char[shaderSource.size() + 1];
		memcpy(fileContents, shaderSource.c_str(), shaderSource.size());
		fileContents[shaderSource.size()] = '\0';

		// Write the result to the output parameters.
		*Data = (LPCVOID)fileContents;
		*Bytes = shaderSource.length();

		return D3D_OK;
	}

	STDMETHOD(Close)(LPCVOID Data)
	{
		delete Data;
		return D3D_OK;
	}
};

Shader* Shader::Compile(
		const ShaderType* shaderType, 
		const VertexFactoryType* vertexFactoryType, 
		const ShaderDefineMap* defineMap)
{
	std::string shaderSource;
	if (!Utility::LoadFileToString(SHADER_PATH + shaderType->SourceFile, shaderSource))
		return NULL;

	LPD3DXBUFFER pShader = NULL;
	LPD3DXBUFFER pErrorMsgs = NULL;
	LPD3DXCONSTANTTABLE pConstantTable = NULL;
	D3DIncludeEnvironment includeEnvironment;

	// build the defines.
	std::vector<D3DXMACRO> defines;
	D3DXMACRO macro = { NULL, NULL };
	for (ShaderDefineMap::const_iterator it(defineMap->begin()), end(defineMap->end()); it != end; it++)
	{
		macro.Name = it->first.c_str();
		macro.Definition = it->second.c_str();
		defines.push_back(macro);
	}
	macro.Name = NULL;
	macro.Definition = NULL;
	defines.push_back(macro);

	// compile the shader base on the source.
	HRESULT hr = D3DXCompileShader(
		shaderSource.c_str(),
		shaderSource.size(),
		&defines[0],
		&includeEnvironment,
		shaderType->Entry.c_str(),
		shaderType->Target.c_str(),
		0,
		&pShader,
		&pErrorMsgs,  
		&pConstantTable);

	if (FAILED(hr))
	{
		std::cout << "Error: " << static_cast<const char*>(pErrorMsgs->GetBufferPointer()) << std::endl;
		SAFE_RELEASE(pErrorMsgs);
		return NULL;
	}
	
	// build constant table.
	D3DXCONSTANTTABLE_DESC tableDesc;
	hr = pConstantTable->GetDesc(&tableDesc);
	if (FAILED(hr))
		return NULL;

	// parse the contant table.
	NameIndexMap paramNameIndexMap;
	for (UINT i = 0; i < tableDesc.Constants; i++)
	{
		D3DXHANDLE hConstant = pConstantTable->GetConstant(NULL, i);        
		
        UINT numParams = 1;
        D3DXCONSTANT_DESC constDesc;
		hr = pConstantTable->GetConstantDesc(hConstant, &constDesc, &numParams);
		if (FAILED(hr))
			continue;

		std::string paramName(constDesc.Name);
		if (paramName.at(0) == '$')
			paramName.erase(paramName.begin());
		paramNameIndexMap[paramName] = constDesc.RegisterIndex;
	}

	SAFE_RELEASE(pConstantTable);

	IDirect3DVertexShader9* pVertexShader = NULL;
	IDirect3DPixelShader9* pPixelShader = NULL;
	LPDIRECT3DDEVICE9 pDevice = Render::Instance().getDevice();

	// create shader base on the shader type.
	if (shaderType->Type == ShaderType::EST_Vertex)
		pDevice->CreateVertexShader((DWORD*)pShader->GetBufferPointer(), &pVertexShader);
	else if (shaderType->Type == ShaderType::EST_Pixel)
		pDevice->CreatePixelShader((DWORD*)pShader->GetBufferPointer(), &pPixelShader);

	// there is no shader created.
	if (pVertexShader == NULL && pPixelShader == NULL)
		return NULL;

	// done with the shader.
	Shader* shader = shaderType->Create();
	shader->mVertexShader = pVertexShader;
	shader->mPixelShader = pPixelShader;
	shader->mParamNameIndexMap = paramNameIndexMap;
	shader->mVertexFactoryType = vertexFactoryType;
	if (!shader->Init())
	{
		delete shader;
		return NULL;
	}

	return shader;
}
